import os
import json
import time
from flask import Flask, request, jsonify, Response
from openai import OpenAI
import dashscope
import uuid

# ===================== 基础配置 =====================
app = Flask(__name__)
app.json.ensure_ascii = False  # 返回 JSON 中文不转义

UPLOAD_FOLDER = os.environ.get("UPLOAD_FOLDER", "uploads")
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# DashScope 配置
dashscope.api_key = 'sk-b42f646a808549e099932167d32f2a9c'
DEFAULT_ASR_MODEL = 'qwen-audio-asr'

# Qwen-Max 客户端
client = OpenAI(
    api_key='sk-b42f646a808549e099932167d32f2a9c',
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

ALLOWED_EXTS = {'.wav', '.mp3', '.m4a', '.flac', '.aac', '.ogg', '.opus'}

# ===================== 工具函数 =====================
def allow_file(filename: str) -> bool:
    return os.path.splitext(filename)[1].lower() in ALLOWED_EXTS

def real_asr(file_path):
    """调用 DashScope ASR，将本地音频转文字"""
    audio_uri = f"file://{os.path.abspath(file_path)}"
    try:
        response = dashscope.MultiModalConversation.call(
            model=DEFAULT_ASR_MODEL,
            messages=[{"role": "user", "content": [{"audio": audio_uri}]}],
            result_format="message"
        )
        text = response['output']['choices'][0]['message']['content'][0]['text']
        return text
    except Exception as e:
        return f"ASR 解析失败: {e}"

# ===================== 路由 =====================
@app.route("/health", methods=["GET"])
def health():
    return jsonify({"status": "ok", "service": "asr+qwen"})


@app.route("/infer", methods=["POST"])
def infer():
    try:
        audio_files = request.files.getlist("audio")
        json_file = request.files.get("json")
        if not audio_files or not json_file:
            return jsonify({"error": "缺少 audio 或 json 参数"}), 400

        # 读取 JSON 内容
        user_json = json.load(json_file)

        # 音频转写
        asr_texts = []
        for audio in audio_files:
            if not allow_file(audio.filename):
                return jsonify({"error": f"文件类型不支持: {audio.filename}"}), 400
            save_path = os.path.join(UPLOAD_FOLDER, f"{uuid.uuid4().hex}_{audio.filename}")
            audio.save(save_path)
            text = real_asr(save_path)
            asr_texts.append(text)
            # 清理临时文件
            os.remove(save_path)

        combined_asr_text = "\n".join(asr_texts)

        # 系统提示
        system_prompt = (
            "你是一名专业律师，具有丰富的案件分析和法律意见撰写经验。"
            "请结合以下信息为用户提供详细可操作的劳动法意见，"
            "包括分析理由和法律依据，输出内容应直接可读、逻辑清晰。"
        )
        print("======================= user_json (上传的 JSON) =======================")
        print(json.dumps(user_json, ensure_ascii=False, indent=2))
        print("======================= combined_asr_text (合并后的音频转写) =======================")
        print(combined_asr_text)
        
        user_message = (
            f"JSON 聊天记录:\n{json.dumps(user_json, ensure_ascii=False)}\n\n"
            f"音频转写内容:\n{combined_asr_text}"
        )

        # 调用 Qwen-Max
        completion = client.chat.completions.create(
            model="qwen-plus",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_message},
            ],
        )

        reply_text = completion.choices[0].message.content

        return reply_text  # 直接返回中文文本

    except Exception as e:
        return jsonify({"error": f"推理失败: {e}"}), 500


# ===================== 启动 =====================
if __name__ == "__main__":
    app.run(host="0.0.0.0", port=6006, debug=True)
